'''
Created on Nov 30, 2010

@author: oabalbin
'''
import os
import subprocess

NODE_MEMORY = 45000.0
NODE_PROCESSORS = 12
MEM_PER_PROCESSOR = int(float(NODE_MEMORY) / NODE_PROCESSORS)

def qsub(job_name, cmd, num_processors, cwd=None, walltime="60:00:00", pmem=None, deps=None, stdout=None, 
         email_addresses=None):
    '''
    job_name: string name of job
    cmd: string containing command-line arguments for job
    num_processors: number of processors to submit with (cannot be greater than the number of processors per node)
    cwd: the "working directory" of the job (allows scripts to access files using relative pathnames)
    walltime: the walltime passed to qsub
    pmem: amount of memory allocated to this job
    deps: 'None' if no dependencies, of a python list of job ids
    stdout: string filename for storing stdout
    email_addresses: list of email addresses to send job information
    '''
    if isinstance(deps, basestring):
        deps = [deps]
    if isinstance(email_addresses, basestring):
        email_addresses = [email_addresses]        
    if cwd is None:
        cwd = os.getcwd()
    num_processors = min(NODE_PROCESSORS, num_processors)
    if pmem is None:
        pmem = MEM_PER_PROCESSOR
    if stdout is None:
        stdout = "${PBS_JOBID}"
    lines = ["#!/bin/sh",
             "#PBS -N %s" % job_name,
             "#PBS -l nodes=%d:ppn=%d,walltime=%s,pmem=%dmb" % (1, num_processors, walltime, pmem),
             "#PBS -l qos=arul_flux",
             "#PBS -A arul_flux",
             "#PBS -q flux",
             "#PBS -V",
             "#PBS -j oe",
             "#PBS -o %s/%s" % (cwd, stdout)]
    if email_addresses is not None:        
        lines.append("#PBS -m bae")
        for email_address in email_addresses:
            lines.append("#PBS -M %s" % email_address)
    if deps is not None:
        lines.append("#PBS -W depend=afterok:%s" % (":".join([d for d in deps])))    
    lines.extend(["cat $PBS_NODEFILE", 
                  "cd %s" % (cwd), 
                  cmd])
    print '\n'.join(lines)
    p = subprocess.Popen("qsub", stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    p.stdin.write('\n'.join(lines))
    job_id = p.communicate()[0]
    #job_id = 'myjob'
    return job_id.strip()


# Depracated
def write_pbs_header(qsubfile, jname, nnodes, ppn, pmem, wt):
    
    qsubfile.write('#!/bin/sh\n')
    qsubfile.write('#PBS -S /bin/sh\n')
    qsubfile.write('#PBS -N '+str(jname)+'\n')
    qsubfile.write('#PBS -l '+'nodes='+str(nnodes)+':ppn='+str(ppn)+',walltime='+wt+'\n')
    qsubfile.write('#PBS -q flux\n')
    qsubfile.write('#PBS -M oabalbin@umich.edu\n')
    qsubfile.write('#PBS -m abe\n')
    qsubfile.write('#PBS -V\n')
    qsubfile.write('#PBS -l qos=arul_flux\n')
    qsubfile.write('#PBS -A arul_flux\n')
    qsubfile.write('# \n')
    qsubfile.write('cat $PBS_NODEFILE \n')
    
    
def pbs_header_list(jname, nnodes, ppn, pmem, wt, jobid):
    '''
    '''
    p1=['#PBS -N '+str(jname)+'\n',
        '#PBS -l '+'nodes='+str(nnodes)+':ppn='+str(ppn)+',pmem='+str(pmem)+',walltime='+wt+'\n',
        '#PBS -q flux\n',
        '#PBS -M oabalbin@umich.edu\n',
        '#PBS -m abe\n']
        
    p3=['#PBS -V\n',
        '#PBS -l qos=arul_flux\n',
        '#PBS -A arul_flux\n',
        '# \n']
    
    if jobid !=-1:
        p2=['#PBS -W depend=afterok:'+str(jobid)+'\n']
        header=p1+p2+p3
    else:
        header=p1+p3
        
    header = [a.replace(',',';') for a in header]
    header = ",".join(header).replace(',',' ').replace(';',',')

    #'cat $PBS_NODEFILE \n'
    return header